#!/usr/bin/env python

# If run in vscode, change the path to site-packages in anaconda env
# import sys
# sys.path.append('/home/<your_workspace_name>/anaconda3/lib/python3.8/site-packages')

import numpy as np
import open3d as o3d
import datetime
# import SVD function
from numpy import linalg as la

def icp_core(points1, points2):
    """
    solve transformation from points1 to points2, points of the same index are well matched
    :param points1: numpy array of points1, size = nx3, n is num of point
    :param points2: numpy array of points2, size = nx3, n is num of point
    :return: transformation matrix T, size = 4x4
    """
    assert points1.shape == points2.shape, 'point could size not match'

    # Initialize transformation matrix T
    T = np.zeros(shape=(4, 4))
    T[0:3, 0:3] = np.eye(3)
    T[3, 3] = 1

    # Todo: step1: calculate centroid
    centroid1 = np.mean(points1, axis=0)   # get mean x,y,z value of points1
    centroid2 = np.mean(points2, axis=0)   # get mean x,y,z value of points2

    # Todo: step2: de-centroid of points1 and points2
    p1i = (points1 - centroid1).T          # subtract centroids by points1
    p2i = (points2 - centroid2).T          # subtract centroids by points2

    # Todo: step3: compute H, which is sum of p1i'*p2i'^T
    H = np.matmul(p1i, p2i.T)              # H = p1i'*p2i'^T

    # Todo: step4: SVD of H (can use 3rd-part lib), solve R and t
    U,sigma,VT = la.svd(H)                 # U * sigma * V^T = H

    # Todo: step5, combine R and t into transformation matrix T
    R = np.matmul(VT.T,U.T)                # R = V * U^T
    t = centroid2 - np.matmul(R,centroid1) # t = centroid2 - R*centroid1
    print('------------Rotation matrix------------')
    print(R)
    print('------------translation matrix------------')
    print(t)
    # combine R and t into T
    # T = [ R t ]
    #     [ 0 1 ]
    T[0:3,0:3] = R
    T[0:3,3] = t.reshape(1,3)
    return T


def svd_based_icp_matched(points1, points2):

    # icp_core should be finished first
    T = icp_core(points1, points2)
    print('------------transformation matrix------------')
    print(T)

    # Todo: calculate transformed point cloud 1 based on T solved above, and name it pcd1_transformed (same format as point1)
    # pcd1_transformed = points1.copy() # comment this line, this only indicates the size of points_2_nearest should be the same as points1
    N = np.size(points1,0)
    pcd1_transformed = np.zeros(shape=(N, 3))
    for i in range(0,N):
        # turn every point array from [x,y,z] to [x,y,z,1]
        # multiply every point array with T
        # and get x_transformed, y_transformed, z_transformed in first three value
        pcd1_transformed[i] = np.matmul(T, np.r_[points1[i], [1]])[0:3]

    # visualization
    mean_distance = mean_dist(pcd1_transformed, points2)
    print('mean_error= ' + str(mean_distance))
    axis_pcd = o3d.geometry.TriangleMesh.create_coordinate_frame(size=100, origin=[0, 0, 0])
    pcd1 = o3d.geometry.PointCloud()
    pcd1.points = o3d.utility.Vector3dVector(points1)
    pcd2 = o3d.geometry.PointCloud()
    pcd2.points = o3d.utility.Vector3dVector(points2)
    pcd1_tran = o3d.geometry.PointCloud()
    pcd1_tran.points = o3d.utility.Vector3dVector(pcd1_transformed)
    pcd1.paint_uniform_color([1, 0, 0])
    pcd2.paint_uniform_color([0, 1, 0])
    pcd1_tran.paint_uniform_color([0, 0, 1])
    # o3d.visualization.draw_geometries([pcd1, pcd2, pcd1_tran, axis_pcd])
    o3d.visualization.draw_geometries([pcd1, pcd2, axis_pcd])


def svd_based_icp_unmatched(points1, points2, n_iter, threshold):
    points_1 = points1.copy()
    points_2 = points2.copy()
    T_accumulated = np.eye(4)

    axis_pcd = o3d.geometry.TriangleMesh.create_coordinate_frame(size=100, origin=[0, 0, 0])
    pcd2 = o3d.geometry.PointCloud()
    pcd2.points = o3d.utility.Vector3dVector(points_2)
    pcd2.paint_uniform_color([0, 1, 0])
    vis = o3d.visualization.Visualizer()
    vis.create_window()
    vis.add_geometry(axis_pcd)
    vis.add_geometry(pcd2)

    start_time = datetime.datetime.now()

    for i in range(n_iter):

        # Todo: for all point in points_1, find nearest point in points_2, and generate points_2_nearest
        # points_2_nearest = points_1.copy() # comment this line, this only indicates the size of points_2_nearest should be the same as points_1
        N = np.size(points_1, 0)
        points_2_nearest = np.zeros(shape=(N, 3))           # initialize
        for j in range(N):                                  # traverse points1
            dif = points_1[j] - points_2                      # difference between one point in points1 and all points in points2
            dis = np.linalg.norm(dif, axis=1)               # compute the distance array
            points_2_nearest[j] = points_2[np.argmin(dis)]   # find points1[j]'s nearest point in points2, and put it in points_2_nearest


        # solve icp
        T = icp_core(points_1, points_2_nearest)

        # Todo: update accumulated T
        # T_accumulated = ?
        T_accumulated = np.dot(T, T_accumulated)            # update


        print('-----------------------------------------')
        print('iteration = ' + str(i+1))
        print('T = ')
        print(T)
        print('accumulated T = ')
        print(T_accumulated)

        # Todo: update points_1
        points_1_homogeneous = np.append(points_1, np.ones(shape=(N, 1)), axis=1)   # homogeneous coordinate (N, 4)
        points_1_homogeneous = (np.dot(T, points_1_homogeneous.T)).T        # [(4, 4) * (4, N)]' = (N, 4)
        points_1 = points_1_homogeneous[:, 0:3]         # (N, 3)

        test1 = np.array([[181.68852093847812, 105.62510479653113, 377.0],
                 [185.35147623882528, 20.588342356967097, 440.0],
                 [364.0958836102045, -147.03589488302043, 798.0]])

        test1_homogeneous = np.append(test1, np.ones(shape=(3, 1)), axis=1)   # homogeneous coordinate (N, 4)
        test1_homogeneous = (np.dot(T_accumulated, test1_homogeneous.T)).T        # [(4, 4) * (4, N)]' = (N, 4)
        test_1 = test1_homogeneous[:, 0:3]         # (N, 3)
        print("test:  \n", test_1)
        test2 = np.array([[278.54977729886025, 103.86488276045722, 332.0],
                         [296.8702246672814, 17.427487834565884, 392.0],
                         [551.0550148178618, -146.97627434471065, 708.0]])
        mean_distance2 = mean_dist(test_1, test2)
        print('mean_error22222222222= ' + str(mean_distance2))


        mean_distance = mean_dist(points_1, points2)
        print('mean_error= ' + str(mean_distance))

        # visualization
        pcd1_transed = o3d.geometry.PointCloud()
        pcd1_transed.points = o3d.utility.Vector3dVector(points_1)
        pcd1_transed.paint_uniform_color([1, 0, 0])
        vis.add_geometry(pcd1_transed)
        vis.poll_events()
        vis.update_renderer()
        vis.remove_geometry(pcd1_transed)

        if mean_distance < 0.00001 or mean_distance < threshold:
            print('fully converged!')
            break
    end_time = datetime.datetime.now()
    time_difference = (end_time - start_time).total_seconds()
    print('time cost: ' + str(time_difference) + ' s')
    vis.destroy_window()
    # o3d.visualization.draw_geometries([pcd2, pcd1_transed])
    o3d.visualization.draw_geometries([axis_pcd, pcd2, pcd1_transed])


def mean_dist(points1, points2):
    dis_array = []
    for i in range(points1.shape[0]):
        dif = points1[i] - points2[i]
        dis = np.linalg.norm(dif)
        dis_array.append(dis)
    return np.mean(np.array(dis_array))


def main():
    # pcd1 = o3d.io.read_point_cloud('data/bun000.ply')
    # pcd2 = o3d.io.read_point_cloud('data/bun045.ply')
    # points1 = np.array(pcd1.points)
    # points2 = np.array(pcd2.points)

    points1 = np.array([[-402.04504329470853, -42.008889599155445, 644.0], [-159.41183553770367, 44.527853744007885, 353.0],
                        [-144.3482810591662, 74.9940356847071, 318.0], [-50.784702144470614, 66.70545403262737, 267.0],
                        [-117.57706772296903, -59.50801841937174, 440.0], [-95.93681728106618, 3.0958586314355454, 435.0],
                        [-32.83066038606956, 104.32452785324925, 331.0], [-1.1025230804069976, 107.3683021830877, 315.0],
                        [-0.6717676177167592, 60.558876474157834, 576.0], [66.471939702134, 65.87857754380389, 564.0],
                        [69.43006768186827, 5.679886590461228, 601.0], [98.99588056154839, 10.728846548201714, 508.0],
                        [142.22266659002744, 3.1314432134060692, 440.0], [162.7601949942169, -160.1419075999, 440.0]])

    points2 = np.array([[-222.19225457129033, -26.951243522489786, 724.0], [-61.45829959762292, 41.634227575983516, 396.0],
                        [-44.898778014026014, 75.8335817001212, 353.0], [22.94969886463898, 68.55732026161348, 277.0],
                        [0.5394066311762953, -49.544776445700975, 462.0], [21.86459271687851, 4.318981983096142, 457.0],
                        [57.528635843429036, 104.78066018851253, 340.0], [86.46086486328015, 106.16440265815645, 318.0],
                        [144.00811727606458, 60.34860259751146, 574.0], [211.87529287419855, 63.30884579564133, 542.0],
                        [223.8220883204413, -1.2777545028198367, 576.0], [232.36960150849265, 14.40522297166335, 473.0],
                        [250.6502906724238, 1.9275918583141558, 403.0], [272.86115702875895, -160.24255972080465, 397.0]])
    # task 1:
    # svd_based_icp_matched(points1, points2)
    # task 2:
    svd_based_icp_unmatched(points1, points2, n_iter=10, threshold=0.1)



if __name__ == '__main__':
    main()
